package com.wy.controller;

import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@CrossOrigin("*")
public class SseController {
    private final ExecutorService executor = Executors.newSingleThreadExecutor();
    private final AtomicBoolean running = new AtomicBoolean(true);

    @GetMapping(value = "/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamSseMvc() {
        SseEmitter emitter = new SseEmitter();

        executor.execute(() -> {
            try {
                int i = 0;
                while (running.get()) {
                    // 发送事件
                    emitter.send(SseEmitter.event()
                            .name("message")
                            .data("Message " + i));
                    // 每隔2秒发送一次消息
                    Thread.sleep(500);
                    i++;
                }
            } catch (IOException | InterruptedException e) {
                emitter.completeWithError(e);
            } finally {
                emitter.complete();
            }
        });

        // 在客户端断开连接时停止发送事件
        emitter.onCompletion(() -> {
            running.set(false);
            executor.shutdown(); // 关闭ExecutorService
        });
        emitter.onTimeout(() -> {
            running.set(false);
            executor.shutdown(); // 关闭ExecutorService
        });

        return emitter;
    }
}
